import numpy as np
import platgo as pg


class DTLZ6(pg.Problem):

    def __init__(self, M: int = 3) -> None:
        self.name = 'DTLZ6'
        self.type['multi'], self.type['many'], self.type['real'], self.type['large'], self.type['expensive'] = [True] * 5
        self.M = M
        self.D = M + 9
        lb = [0] * self.D
        ub = [1] * self.D
        self.borders = np.array([lb, ub])
        super().__init__()

    def cal_obj(self, pop: pg.Population) -> None:
        decs = pop.decs
        pop.cv = np.zeros((pop.N, self.D))
        XM = decs[:, (self.M - 1):]
        g = np.sum((XM ** 0.1), axis=1, keepdims=True)
        decs[:, 1:(self.M - 1)] = 1 * (1 + 2 * g * decs[:, 1:(self.M-1)]) / (2 + 2 * g)
        ones_matrix = np.ones((pop.N, 1))
        f = np.fliplr(np.cumprod(np.hstack([ones_matrix, np.cos(decs[:, :self.M - 1] * np.pi / 2)]), axis=1)) * \
            np.hstack([ones_matrix, np.sin(decs[:, range(self.M - 2, -1, -1)] * np.pi / 2)]) * np.tile(1 + g,
                                                                                                       (1, self.M))
        pop.objv = f

    def get_optimal(self) -> np.ndarray:
        # 目标空间均匀采样函数还未完成
        raise NotImplementedError("get optimal has not been implemented")


if __name__ == '__main__':
    d = DTLZ6()
    pop = d.init_pop(10)

    # pop = pg.Population(decs=np.random.uniform(0, 1, (10, 7)))
    print(d.borders)
    print(pop.decs)
    print(pop.cv)
    d.cal_obj(pop)  # 计算目标函数值
    print(pop.objv)
